'''
:@Author: junpeng_chen
:@Date: 7/25/2023, 1:32:17 PM
:@LastEditors: junpeng_chen
:@LastEditTime: 7/25/2023, 1:32:17 PM
:Description: 数值字面量转换
:Copyright: Copyright (©)}) 2023 junpeng_chen. All rights reserved.
'''
import sys
import os

# 获取父目录的绝对路径
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
# 将父目录添加到sys.path
sys.path.append(parent_dir)
import ast
import re
from utils.log import logger
from utils.chatglm_utils import ChatGLMUtils
from transformers import AutoTokenizer, AutoModel


class MeasureRecognition:
    def __init__(self, model) -> None:
        self.model = model
        self.prompt = "假设你是数值格式转换器，输入字面量，需要输出格式化后的数值。格式要求：如果不存在数值则只生成none即可；单个数值点直接转换为数字；多个数值点则转换为list列表形式；数值范围则识别左边界点和右边界点，其中若某一边无边界点则用none标识无边界，最后将边界点用英文括号括起来，英文小括号表示不包括边界点，英文大括号表示包括边界点；输入：{}\n输出： "

    def filter(self, query, literals,agg):
        """
        过滤子任务1输出的格式化数据，只提取和数值条件相关的数据
        :param query: 自然语言查询
        :param literals: 子任务1-literal&opt识别输出的格式化数据
        :param agg: header对应的聚合函数
        :return: 经过过滤后的只包含数值条件字面量的子任务1-literal&opt识别输出的格式化数据以及相应的聚合函数列表
        """
        measure_literals = []
        measure_agg = []
        for header,func in zip(literals,agg):
            if header["literal"] == "none" or header["timeLevel"] != "":
                # 不存在条件字面量和属于时间维度的header不在考虑范围
                continue
            else:
                # 考虑存在字面量且其中包含数值的
                number_pattern = r"零|一|二|三|四|五|六|七|八|九|十|百|千|万|亿|0|1|2|3|4|5|6|7|8|9"
                m = re.search(number_pattern, header["literal"])
                if m:
                    measure_literals.append(header)
                    measure_agg.append(func)

        return measure_literals,measure_agg

    def has_only_one_pair_of_brackets(self, s):
        """
        判断字符串中只存在一对括号
        :param s: 只存在英文字符的字符串
        :return: 0或1或-1——没有括号，1对括号，其他非法情况
        """
        left_bracket_count = 0
        right_bracket_count = 0

        for char in s:
            if char in ['(', '[', '{']:
                left_bracket_count += 1
            elif char in [')', ']', '}']:
                right_bracket_count += 1

                # 提前判断，如果发现超过1对括号，则可以直接返回False
            if left_bracket_count > 1 or right_bracket_count > 1:
                return -1
        if left_bracket_count == right_bracket_count and left_bracket_count == 1 and bool(
                re.match(r"^[^\(\[]*[\({\[][^()]*[\)}\]][^\(\[]*$", s)) == False:
            # 判断左括号在右括号的左边
            return -1

        return left_bracket_count if left_bracket_count == right_bracket_count else -1

    def is_charset_legal(self, prediciton):
        """
        判断字符串只包含数字或英文逗号或英文括号或一个none
        :param prediciton: 模型预测的有效字符串
        :return: bool值
        """
        pattern = r'^([0-9,()\[\]{}none]+)$'  # 正则表达式模式，用于匹配只包含数字、英文逗号、英文括号或 "none" 子字符串的字符串
        return bool(re.match(pattern, prediciton))

    def find_first_number(self, prediction):
        """
        返回字符串中的第一个数字
        :return: 整数
        """
        try:
            pattern = r'\d+'  # 正则表达式模式，匹配至少一个数字
            match = re.search(pattern, prediction)
            if match:
                return int(match.group())  # 将匹配的字符串转换为整数并返回
            else:
                return None
        except Exception as err:
            logger.error("find first number error[{}]: {}".format(type(err), str(err)))
            raise

    def range_transform(self, range):
        """
        将范围类型的预测转换为合法结果
        :param range: 模型预测的合法范围值，形如(x,y),{x,y},{x,y),(x,y}
        :return: 格式化后的操作符和条件字符串元组，以列表形式返回: [(操作符1, 条件值1),(操作符2, 条件值2)]
        """
        try:
            res = []
            range_patter = r"[\({](.*?),(.*?)[\)}]"
            match = re.search(range_patter, range)
            if match:
                x_str = match.group(1)
                y_str = match.group(2)
                opt_x = '>' if '(' in range else ">="
                opt_y = '<' if ')' in range else "<="
                if x_str != 'none':
                    res.append((opt_x, x_str))
                if y_str != 'none':
                    res.append((opt_y, y_str))

            return res
        except Exception as err:
            logger.error("range transform error[{}]: {}".format(type(err), str(err)))
            raise

    def formatter(self, prediciton: str):
        """
        对模型预测的非none输出进行一定地格式化并返回
        :param prediciton: 模型预测的有效字符串
        :return: 格式化后的操作符和条件字符串元组，以列表形式返回: [(操作符1, 条件值1),(操作符2, 条件值2)]
        """
        res = []
        try:
            '''
            ！！问题：单位的事情还没有解决，如“半年”可能模型识别的是0.5，“50%”模型识别的可能是0.5或50%
            0. 基础的字符串前处理/判断：literal应只包含数值、英文逗号、英文括号（()、{}、[]）和括号中的none，且括号若存在且恒为1对
            1. 不包括英文小括号和英文大括号的单个数值和列表不需要转换，直接作为条件值返回
            2. 括号若不成对出现则根据数值的个数进行删除或替换为列表形式返回
            3. 范围值的转换（()表示不包括边界点，{}表示包括边界点，括号的区别只影响操作符，下面只以英文小括号为例，x为任意整数数值）：
                3.1. (x,none)将x作为条件值，操作符为 >
                3.2. (none,x)将x作为条件值，操作符为 <
                3.3. (x,y)将拆分为两个condition，其中一个为('>','x')，另一个为('<',y)
            '''
            # 前处理/判断
            # 替换非法字符
            literal = prediciton.replace(" ", "")
            literal = literal.replace("，", ",")
            literal = literal.replace("（", "(")
            literal = literal.replace("【", "[")
            literal = literal.replace("）", ")")
            literal = literal.replace("】", "]")

            # 括号对数判断——只能由一个括号或没有括号
            bracket_num = self.has_only_one_pair_of_brackets(literal)
            if bracket_num == -1:
                return res
            elif bracket_num == 1:
                # 前提：只有一对括号
                # 删除左括号之前的字符和删除有括号后面的字符
                left_brackets_pattern = r'([{\[\(])'
                right_brackets_pattern = r'([}\]\)])'
                left_matches = re.finditer(left_brackets_pattern, literal)
                right_matches = re.finditer(right_brackets_pattern, literal)
                left_indices = [match.start() for match in left_matches]
                right_indices = [match.start() for match in right_matches]
                literal = literal[left_indices[0]:right_indices[0] + 1]
                try:
                    if type(ast.literal_eval(literal)) == list:
                        # 如果是列表
                        res.append(("in", str(literal)))
                except:
                    # 格式化
                    # 确保字面量只存在()/{}在内的一对英文括号
                    literal = literal.replace("[", '(')
                    literal = literal.replace("]", ')')
                    res.extend(self.range_transform(literal))
            else:
                # 没有括号，说明是数值
                # 解决方案——提取发现的第一个数字进行返回
                num = self.find_first_number(literal)
                if num:
                    res.append(('==', str(num)))
            return res
        except Exception as err:
            logger.error("measure prediction formatting error[{}]: {}".format(type(err), str(err)))
            raise

    def transform(self, query, measure_literals, measure, measure_agg):
        """
        将字面量转换为对应的数值条件
        :param query: 自然语言查询
        :param measure_literals: 经过过滤后的只包含数值条件字面量的子任务1-literal&opt识别输出的格式化数据
        :param measure: 列表形式，其中元素为类型为“指标”的列名id
        :param measure_agg: 列表形式，其中元素为measure_header对应的聚合函数
        :return: conds中的数字条件部分
        """
        conds = []
        having = []
        origin_output = {}
        for header,agg in zip(measure_literals,measure_agg):
            # 规则前处理
            literal = header['literal']

            # 模型预测
            try:
                prediction = self.model.predict(self.prompt.format(str(literal)))
            except Exception as err:
                logger.error("measure prediction error happened![{}]: {}".format(type(err), str(err)))
                raise

            # 规则后处理
            if prediction == 'none':
                # 说明不属于数值类型条件
                continue
            else:
                origin_output[str(literal)] = prediction
                # 对有效预测输出进行格式化
                res = self.formatter(prediction)
                # measure模型预测的内容也隐含了操作符，暂时以字面量模型识别的操作符为准（因为暂时还没想好用什么指标来评判）
                for item in res:
                    operator = item[0] if header["operator"] in ['between', 'none'] else header["operator"]
                    number = ""
                    if item[1].isdigit():
                        number = float(item[1])
                    if header['id'] not in measure:
                        conds.append([header["id"], operator, item[1]])
                    else:
                        having.append([header["id"], operator,agg, number])

        return {"origin":origin_output,'processed': (conds, having)}


# 测试
# model = ChatGLMUtils()
# model.load_model(local_loading=True,
#                  model_path=f"/home/cike/amer/smartbi/值识别_联调/val-reg-with-glm/models/pytorch_model.bin",cuda_index=4)

# if __name__ == "__main__":
#     measure_reg = MeasureRecognition(model)
#     query = "今年合同金额大于五百万的分部有哪些？"
#     literals = [{'member_amount': 0, 'name': '年', 'literal': '2023', 'operator': '==', 'timeLevel': 'year',
#                  'id': 'AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.02de0a83c07e6c06b5f6dc8d09f9fce7-C_DATE_Year-LEVEL-1648634671682',
#                  'type': 'STRING'},
#                 {'member_amount': 23, 'name': '销售分部', 'literal': 'none', 'operator': 'none', 'timeLevel': '',
#                  'id': 'AUGMENTED_DATASET_FIELD.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.2b4ff0f606441075f12c1be466e3642e',
#                  'type': 'STRING'},
#                 {'member_amount': 0, 'name': '合同金额', 'literal': '大于五百万', 'operator': '>', 'timeLevel': '',
#                  'id': 'AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8a8ae5ca0178549554951b9501785d983aaa005e',
#                  'type': 'DOUBLE'}]
#     measure = ['AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8a8ae5ca0178549554951b9501785d983aaa005e']
#     measure_literals  = measure_reg.filter(query, literals)
#     print(measure_literals)
#     print(measure_reg.transform(query, measure_literals, measure))
